import sys

from WithSecure.common import cli, console

from drozer import payload
from drozer.exploit import builder
from drozer.util import StoreZeroOrTwo

class ExploitManager(cli.Base):
    """
    drozer exploit COMMAND [OPTIONS]
    
    A utility for crafting exploits to deliver drozer onto an Android device.
    """
    
    def __init__(self):
        cli.Base.__init__(self)
        
        self.builder = builder.Builder()
    
    def args_for_build(self):
        self._parser.add_argument("exploit", help="specify the exploit module to use")
        self._parser.add_argument("--credentials", default=None, nargs=2, metavar=("username", "password"), help="add a username/password pair that can be used to upload files to the server")
        self._parser.add_argument("--payload", default=None, help="specify the payload module to use")
        self._parser.add_argument("--permission", "-p", nargs="+", help="add permissions to the manifest, where a drozer Agent can be deployed")
        self._parser.add_argument("--server", default=None, metavar="HOST[:PORT]", help="specify the address and port of the drozer server")
        self._parser.add_argument("--push-server", default=None, metavar="HOST[:PORT]", help="specify a different drozer server to upload the exploit resources to")
        self._parser.add_argument("--ssl", action=StoreZeroOrTwo, help="enable SSL, optionally specifying the key and certificate", nargs="*")
        
    def do_build(self, arguments):
        """build an exploit"""
        
        self.builder.build(arguments.exploit, arguments)

    def args_for_info(self):
        self._parser.add_argument("module", help="specify the exploit module to use")
        
    def do_info(self, arguments):
        """prints information about an exploit module"""
        
        module = self.builder.module(arguments.module)
        
        print(module.usage.formatted_description())
        
    def do_list(self, arguments):
        """list the available exploit modules"""
        
        sys.stdout.write(console.format_dict(dict(map(lambda m: [m, self.builder.module(m).name], self.builder.modules()))) + "\n")
    
    def before_parse_args(self, argv):
        """
        Allow exploit and payload modules to add additional arguments.
        """
        args, unknown = self.parse_known_args(self._parser, argv)
        
        if hasattr(args, 'exploit') and args.exploit != None and args.exploit != "":
            try:
                self.builder.module(args.exploit).add_arguments(self._parser)
            except KeyError:
                pass
        if hasattr(args, 'payload') and args.payload != None and args.payload != "" and args.payload != "none":
            try:
                self.builder.module(args.payload).add_arguments(self._parser)
            except KeyError:
                pass
        
    def get_completion_suggestions(self, action, text, line, **kwargs):
        if action.dest == "exploit":
            return self.builder.modules()
        elif action.dest == "payload":
            target_exploit = line.split(" ")[3]
            return self.builder.modules(exploit=target_exploit,module_type="payload")
        elif action.dest == "resource":
            return ["/view.jsp"]
        elif action.dest == "server" or action.dest == "push_server":
            return ["localhost:31415"]
            
